{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# **Anomaly Detection for spirit dataset**\n"
      ],
      "metadata": {
        "id": "NxcscygxQKNS"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WkkgktabSyKZ"
      },
      "source": [
        "# Import libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "p-_F7uwWUcCf",
        "outputId": "b274c899-f9f2-4a95-a86e-cfba28f95283"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[nltk_data] Downloading package stopwords to /root/nltk_data...\n",
            "[nltk_data]   Unzipping corpora/stopwords.zip.\n",
            "[nltk_data] Downloading package punkt to /root/nltk_data...\n",
            "[nltk_data]   Unzipping tokenizers/punkt.zip.\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import math, copy, time\n",
        "import pickle \n",
        "import os\n",
        "import re\n",
        "\n",
        "import signal\n",
        "import nltk\n",
        "from nltk import word_tokenize\n",
        "from nltk.corpus import stopwords\n",
        "\n",
        "nltk.download('stopwords')\n",
        "nltk.download('punkt')\n",
        "\n",
        "from tqdm import trange\n",
        "\n",
        "import random\n",
        "import re\n",
        "\n",
        "import tensorflow.keras.backend as K\n",
        "from tensorflow.keras.layers import Layer\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras.callbacks import Callback\n",
        "import keras\n",
        "\n",
        "np.random.seed(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IDM8RiW1Sc-N"
      },
      "source": [
        "# First enable GPU parallel execution"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "mirrored_strategy = tf.distribute.MirroredStrategy()"
      ],
      "metadata": {
        "id": "BCeoG8M467sp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oYtNm7lES3Ng"
      },
      "source": [
        "# Mount dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ME281tJIUOXa",
        "outputId": "9243246d-23c8-40f6-c194-2351966bca07"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "# special to google colab | put your 4 Datasets into this path \n",
        "folder_path = '/content/drive/MyDrive/LogSy'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpcmiFCpS6xQ"
      },
      "source": [
        "# Read in the datasets using preprocessor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5b8wEE-yrWCB"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "Standard tokenizer + do what the paper says\n",
        "\"\"\"\n",
        "\"\"\"\n",
        "Do not run this script in Windows, it will not work due to use of signal\n",
        "\n",
        "We also use authors preprocessor to evaluate our model just in case we made something wrong here\n",
        "\"\"\"\n",
        "\n",
        "class TimeoutException(Exception):  # Custom exception class\n",
        "    pass\n",
        "\n",
        "\n",
        "def timeout_handler(signum, frame):  # Custom signal handler\n",
        "    raise TimeoutException\n",
        "\n",
        "\n",
        "# Change the behavior of SIGALRM\n",
        "signal.signal(signal.SIGALRM, timeout_handler)\n",
        "\n",
        "\n",
        "class DataTokenizer:\n",
        "    def __init__(self):\n",
        "        self.word2index = {'[PAD]': 0, '[CLS]': 1, '[MASK]': 2}\n",
        "        self.num_words = 3\n",
        "        self.stop_words = set(stopwords.words('english'))\n",
        "\n",
        "    def tokenize(self, message):\n",
        "        # paper section IV: Tokenization processing\n",
        "        message = message.lower()\n",
        "        message = re.sub(r'/.*:', '', message, flags=re.MULTILINE)  # filter for endpoints\n",
        "        message = re.sub(r'/.*', '', message, flags=re.MULTILINE)\n",
        "        message = word_tokenize(message)  # remove non words\n",
        "        message = (word for word in message if word.isalpha())  # generator  # remove numerical\n",
        "        message = [word for word in message if word not in self.stop_words]  # remove nltk common stopwords\n",
        "        message = ['[CLS]'] + message  # add embedding token\n",
        "        for word_idx, word in enumerate(message):  # convert to value\n",
        "            if word not in self.word2index:\n",
        "                self.word2index[word] = self.num_words\n",
        "                self.num_words += 1\n",
        "            message[word_idx] = self.word2index[word]\n",
        "        return message\n",
        "tokenizer = DataTokenizer()\n",
        "\n",
        "from tqdm.auto import tqdm\n",
        "class DataImporter:\n",
        "    \"\"\"\n",
        "    loads data set from the raw dataset\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, log_template, dataset_folder_path, dataset_name, dataset_step=1,\n",
        "                 dataset_limit=100000, dataset_type='main', normal_indicator: str = '-', aux_count=50000,\n",
        "                 chunk: bool = True):\n",
        "        self.log_template = log_template  # a template containing <Token{n}> and <Message>\n",
        "        self.log_dataframe = None\n",
        "        self.dataset_folder_path: str = dataset_folder_path  # path to the dataset folder\n",
        "        self.dataset_name: str = dataset_name  # full name of raw dataset\n",
        "        self.step: int = dataset_step  # step taken to sample auxiliary dataset\n",
        "        self.log_template_regex: re = re.compile(r'')\n",
        "        self.log_template_headers: list[str] = []\n",
        "        self.limit: int = dataset_limit  # used for faster experiment only\n",
        "        self.dataset_type: str = dataset_type\n",
        "        self.normal_indicator: str = normal_indicator  # a sign indicating the log line is anomaly\n",
        "        self.aux_count: int = aux_count\n",
        "        self.chunk: bool = chunk\n",
        "\n",
        "    def log_loader(self):\n",
        "        \"\"\"\n",
        "        read from IO stream and only take the actual log message based on template\n",
        "        :return:\n",
        "        \"\"\"\n",
        "        log_messages = []\n",
        "        true_labels = []\n",
        "        with open(os.path.join(self.dataset_folder_path, self.dataset_name), 'r', encoding=\"latin-1\") as ds:\n",
        "            for line_no, line in enumerate(tqdm(ds, miniters=1)):\n",
        "                if line_no != 1 and line_no % 20000000 == 1:  # 15 file chunks\n",
        "                    with open(f'dataset/tbird/{self.dataset_name}_chunked_msg_line{line_no}', 'wb') as message_file:\n",
        "                        # print(log_messages)\n",
        "                        # log_messages_array = np.asanyarray(log_messages).reshape(-1, 1)\n",
        "                        print(len(log_messages))\n",
        "                        l = len(log_messages)\n",
        "                        tokenized = [tokenizer.tokenize(log_message) for log_message in\n",
        "                                     tqdm(log_messages, position=0, leave=True, total=l)]\n",
        "                        tokenized_np = np.asanyarray(tokenized)\n",
        "                        del tokenized\n",
        "                        pickle.dump(tokenized_np, message_file)\n",
        "                        log_messages = []  # reset\n",
        "                # Start the timer. Once 10 seconds are over, a SIGALRM signal is sent.\n",
        "                try:\n",
        "                    try:\n",
        "                        signal.alarm(30)\n",
        "\n",
        "                        match = self.log_template_regex.search(line.strip())\n",
        "\n",
        "                        if not match:\n",
        "                            continue\n",
        "                        label_decider = lambda x: 0 if x == self.normal_indicator else 1\n",
        "                        true_labels.append(label_decider(match.group('Token0')))\n",
        "                        # message = tokenizer.tokenize(match.group('Message'))\n",
        "                        # log_messages.append(message)\n",
        "                        log_messages.append(match.group('Message'))\n",
        "                        # print('message after tokenize ', message, log_messages)\n",
        "                        # print(self.log_template_headers)\n",
        "                    except Exception as e:  # noqa\n",
        "                        print(e)\n",
        "                        # print(e) # will skip those without normal indications('-' OR 'warn')\n",
        "                        pass\n",
        "                except TimeoutException:\n",
        "                    print(\"Regex hang detected, skipping\")\n",
        "                    continue  # catastrophic backtracking\n",
        "                else:\n",
        "                    signal.alarm(0)\n",
        "                if line_no == self.limit:\n",
        "                    break\n",
        "        return log_messages, np.array(true_labels)  # remaining log_messages and all of labels\n",
        "\n",
        "    def load(self):\n",
        "        self.log_template_matcher()\n",
        "\n",
        "        self.log_dataframe = self.log_loader()\n",
        "\n",
        "\n",
        "    def log_template_matcher(self):\n",
        "        headers = []\n",
        "        splitters = re.split(r'(<[^<>]+>)', self.log_template)\n",
        "        regex = ''\n",
        "        for k in range(len(splitters)):\n",
        "            if k % 2 == 0:\n",
        "                splitter = re.sub(' +', '\\\\\\s+', splitters[k])\n",
        "                regex += splitter\n",
        "            else:\n",
        "                header = splitters[k].strip('<').strip('>')\n",
        "                regex += '(?P<%s>.+?)' % header\n",
        "                headers.append(header)\n",
        "        print(regex)\n",
        "        regex = re.compile('^' + regex + '$')\n",
        "\n",
        "        self.log_template_headers, self.log_template_regex = headers, regex"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QmysvQ57IyFV"
      },
      "outputs": [],
      "source": [
        "# original authors' data processor\n",
        "\n",
        "class LogReader:\n",
        "    def __init__(self, log_format, log_name, indir='./', outdir='./result/', rex=[], every_n=10, max_lines=2000000):\n",
        "    \n",
        "        self.path = indir\n",
        "       \n",
        "        self.logName = log_name\n",
        "        self.savePath = outdir\n",
        "        self.df_log = None\n",
        "        self.log_format = log_format\n",
        "        self.rex = rex\n",
        "        self.every_n = every_n\n",
        "        self.max_lines = max_lines\n",
        "    def log_to_dataframe(self, log_file, regex, headers, logformat):\n",
        "            \"\"\" Function to transform log file to dataframe \n",
        "            \"\"\"\n",
        "            log_messages = []\n",
        "            linecount = 0\n",
        "            \n",
        "            if self.max_lines:\n",
        "                with open(log_file, 'r', encoding=\"latin-1\") as fin:\n",
        "                    for i,  line in enumerate(fin):\n",
        "                        if i % self.every_n == 0:\n",
        "                            try:\n",
        "                                match = regex.search(line.strip())\n",
        "                                message = [match.group(header) for header in headers]\n",
        "                                log_messages.append(message)\n",
        "                                linecount += 1\n",
        "                            except Exception as e:\n",
        "                                pass\n",
        "                        if i==self.max_lines:\n",
        "                            break\n",
        "            else:\n",
        "                with open(log_file, 'r', encoding=\"latin-1\") as fin:\n",
        "                    for i,  line in enumerate(fin):\n",
        "                        if i % self.every_n == 0:\n",
        "                            try:\n",
        "                                match = regex.search(line.strip())\n",
        "                                message = [match.group(header) for header in headers]\n",
        "                                log_messages.append(message)\n",
        "                                linecount += 1\n",
        "                            except Exception as e:\n",
        "                                pass\n",
        "            logdf = pd.DataFrame(log_messages, columns=headers)\n",
        "            logdf.insert(0, 'LineId', None)\n",
        "            logdf['LineId'] = [i + 1 for i in range(linecount)]\n",
        "            return logdf\n",
        "\n",
        "\n",
        "    def generate_logformat_regex(self, logformat):\n",
        "        \"\"\" Function to generate regular expression to split log messages\n",
        "        \"\"\"\n",
        "        headers = []\n",
        "        splitters = re.split(r'(<[^<>]+>)', logformat)\n",
        "        regex = ''\n",
        "        for k in range(len(splitters)):\n",
        "            if k % 2 == 0:\n",
        "                splitter = re.sub(' +', '\\\\\\s+', splitters[k])\n",
        "                regex += splitter\n",
        "            else:\n",
        "                header = splitters[k].strip('<').strip('>')\n",
        "                regex += '(?P<%s>.*?)' % header\n",
        "                headers.append(header)\n",
        "        regex = re.compile('^' + regex + '$')\n",
        "        return headers, regex\n",
        "\n",
        "    def load_data(self):\n",
        "        headers, regex = self.generate_logformat_regex(self.log_format)\n",
        "        self.df_log = self.log_to_dataframe(os.path.join(self.path, self.logName), regex, headers, self.log_format)\n",
        "        \n",
        "        \n",
        "class LogTokenizer:\n",
        "    def __init__(self):\n",
        "        self.word2index = {'[PAD]':0, '[CLS]':1, '[MASK]':2}\n",
        "        self.index2word = {0:'[PAD]', 1:'[CLS]', 2:'[MASK]'}\n",
        "        self.n_words = 3  # Count SOS and EOS\n",
        "        self.stop_words = set(stopwords.words('english'))\n",
        "        self.regextokenizer =  nltk.RegexpTokenizer('\\w+|.|')\n",
        "        \n",
        "    def addWord(self, word):\n",
        "        if word not in self.word2index:\n",
        "            self.word2index[word] = self.n_words\n",
        "            self.index2word[self.n_words] = word\n",
        "            self.n_words += 1\n",
        "\n",
        "    def tokenize(self, sent):\n",
        "        sent = re.sub(r'\\/.*:', '', sent, flags=re.MULTILINE)\n",
        "        sent = re.sub(r'\\/.*', '', sent, flags=re.MULTILINE)\n",
        "        sent = self.regextokenizer.tokenize(sent)\n",
        "        sent = [w.lower() for w in sent]\n",
        "        sent = [word for word in sent if word.isalpha()]\n",
        "        sent = [w for w in sent if not w in self.stop_words]\n",
        "        sent = ['[CLS]'] + sent\n",
        "        for w in range(len(sent)):\n",
        "            self.addWord(sent[w])\n",
        "            sent[w] = self.word2index[sent[w]]\n",
        "        return sent\n",
        "    \n",
        "def get_data(log_file, input_dir, output_dir, log_format, regex=[], every_n=10, aux=0, max_lines=5000000):\n",
        "    reader = LogReader(log_format, log_file, indir=input_dir, outdir=output_dir, rex=regex, every_n=every_n, max_lines=max_lines)\n",
        "    reader.load_data()\n",
        "    log_payload, true_labels = reader.df_log.Content, np.where(reader.df_log.t.values=='-',0,1)\n",
        "    del reader\n",
        "    if aux != 0:\n",
        "        df_anomalies = log_payload.iloc[true_labels.flatten()==1].sample(n=aux).values\n",
        "        df_normal = log_payload.iloc[true_labels.flatten()==0].sample(n=aux).values\n",
        "        return df_normal, df_anomalies\n",
        "    else:\n",
        "        return log_payload, true_labels\n",
        "    \n",
        "def get_data_special(log_file, input_dir, output_dir, log_format, regex=[], every_n=10, aux=0, max_lines=2000000):\n",
        "    reader = LogReader(log_format, log_file, indir=input_dir, outdir=output_dir, rex=regex, every_n=every_n, max_lines=max_lines)\n",
        "    reader.load_data()\n",
        "    log_payload, true_labels = reader.df_log.Content, np.where(reader.df_log.t.values!='FATAL',0,1)\n",
        "    del reader\n",
        "\n",
        "    return log_payload, true_labels\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "# First 3 aux datasets "
      ],
      "metadata": {
        "id": "941425gnirAp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZr6kGpZUcCj"
      },
      "outputs": [],
      "source": [
        "log_file = 'bgl2' # first 5 million, note bgl only has 5m when full dataset\n",
        "input_dir  = folder_path  # The input directory of log file\n",
        "output_dir = '.'  # The output directory of parsing results\n",
        "\n",
        "log_format = '<t> <Timestamp> <Date> <Node> <Time> <NodeRepeat> <Type> <Component> <Level> <Content>'  #BGL\n",
        "\n",
        "aux_normal, aux_anomalies = get_data(log_file, input_dir, output_dir, log_format, every_n=1, aux=200000, max_lines=False) # use 200000 when its tbird2 small (5m)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LfyuPftFUcCj",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "log_file = 'tbird2_small' # small meaning the dataset using only first 5m, note model doesn't converge well when sampling 5m instead of first 5m even if using original authors code\n",
        "input_dir  = folder_path\n",
        "log_format = '<t> <Timestamp> <Date> <User> <Month> <Day> <Time> <Location> <Component>(\\[<PID>\\])?: <Content>'  #thunderbird\n",
        "output_dir = '.'\n",
        "every_n = 1\n",
        "aux_size = 200000\n",
        "aux_normal1, aux_anomalies1= get_data(log_file, input_dir, output_dir, log_format, every_n=1, aux=200000, max_lines=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LdSU3380UcCk"
      },
      "outputs": [],
      "source": [
        "log_file = 'Intrepid_RAS_0901_0908_scrubbed' # this one cannot be used as target, its too obvious\n",
        "input_dir  = folder_path\n",
        "output_dir = '.' \n",
        "log_format = '<f> <a>          <c>       <d>                  <e>    <t> <Content>'  #BGP\n",
        "aux_anomalies_t, tl = get_data_special(log_file, input_dir, output_dir, log_format, every_n=1, aux=False, max_lines=False)\n",
        "aux_normal2, aux_anomalies2  = aux_anomalies_t[tl==1], aux_anomalies_t[tl==0] # reverse this [TODO]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oudsVRz7TL32"
      },
      "source": [
        "## Concat the auxiliary dataasets to one single piece, then sample if it was more than a threshold\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "taWdIBBLUcCl"
      },
      "outputs": [],
      "source": [
        "# concatenate\n",
        "aux_anomalies = np.append(aux_anomalies, aux_anomalies1)\n",
        "aux_anomalies_full = np.append(aux_anomalies, aux_anomalies2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a04tdsT-UcCl"
      },
      "outputs": [],
      "source": [
        "# sample\n",
        "if aux_anomalies.shape[0] > 500000:\n",
        "  aux_anomalies = np.random.choice(aux_anomalies_full, size=500000, replace=False)  # changing size doesn't seem to change results\n",
        "  print(f\"Selected {aux_anomalies.shape} samples out of {len(aux_anomalies_full)} aux samples\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Then target dataset - switch with aux if you need to evaluate different target"
      ],
      "metadata": {
        "id": "beUJ6Ow_PQhD"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1scTBiN7QRSc"
      },
      "outputs": [],
      "source": [
        "#read the datasets\n",
        "log_file = 'spirit_small'\n",
        "input_dir  = folder_path\n",
        "output_dir = '.' \n",
        "log_format = '<t> <Timestamp> <Date> <User> <Month> <Day> <Time> <Location> <Content>'  #spirit2\n",
        "\n",
        "log_payload, true_labels = get_data(log_file, input_dir, output_dir, log_format, every_n=1, aux=False, max_lines=5000000)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eIkTZFuQWuvq"
      },
      "outputs": [],
      "source": [
        "### Never shuffle the data here, or there will be info leak"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H4c-K9T7UcCm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "336ed471-33d1-4946-98fb-89af667f524c"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "4973636"
            ]
          },
          "metadata": {},
          "execution_count": 68
        }
      ],
      "source": [
        "df_size = len(log_payload)\n",
        "df_size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ufdIgBLJUcCn"
      },
      "outputs": [],
      "source": [
        "true_labels = true_labels.reshape(-1,1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4z6oK4-TUcCn",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d47bf034-ce46-4ed3-91a0-1239f4f779cc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "log shape (5373636, 1), log labels shape (5373636,)\n"
          ]
        }
      ],
      "source": [
        "#append the anomalies to the full data\n",
        "log_payload = np.append(log_payload.values.reshape(-1,1), aux_anomalies.reshape(-1,1), axis=0)\n",
        "true_labels = np.append(true_labels, np.ones(len(aux_anomalies)).reshape(-1,1), axis=0).flatten()\n",
        "print(f'log shape {log_payload.shape}, log labels shape {true_labels.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MrQp5QZxXy-Z"
      },
      "source": [
        "# Tokenize!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3O2uU67VUcCo",
        "scrolled": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d204503c-6616-4dc0-9f59-96c41f76fec9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 5373636/5373636 [01:24<00:00, 63461.24it/s]\n"
          ]
        }
      ],
      "source": [
        "tokenizer = LogTokenizer()\n",
        "df_len = int(log_payload.shape[0]) # the num of logs\n",
        "data_tokenized = []\n",
        "for i in trange(df_len):\n",
        "    tokenized = tokenizer.tokenize(log_payload[i][0])\n",
        "    data_tokenized.append(tokenized)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RtWHQxFmRbTK"
      },
      "outputs": [],
      "source": [
        "# pickle to file\n",
        "with open('tokenized','wb') as tokenized_file:\n",
        "  pickle.dump(data_tokenized, tokenized_file)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dTfbilkfUl2q"
      },
      "outputs": [],
      "source": [
        "# save it for multiple executions with different split rate\n",
        "with open('tokenized','rb') as token:\n",
        "  data_tokenized = pickle.load(token)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3cWcsFTUUcCo",
        "outputId": "ac87751d-bfdd-475b-c276-6764edd5b2e1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-74-18a50203c19c>:1: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
            "  data_token_indexed = np.asanyarray(data_tokenized)\n"
          ]
        }
      ],
      "source": [
        "data_token_indexed = np.asanyarray(data_tokenized)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PTrvJUVgaxWl"
      },
      "source": [
        "#Standard transformer encoder with multi-head *attention*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iMyZcAlm9rKW"
      },
      "outputs": [],
      "source": [
        "class TransformerEncoder(layers.Layer): # standard tensformer encoder\n",
        "    def __init__(self, embed_dim, num_heads, dense_dim, **kwargs):\n",
        "        super(TransformerEncoder, self).__init__(**kwargs)\n",
        "        self.embed_dim = embed_dim\n",
        "        self.dense_dim = dense_dim\n",
        "        self.num_heads = num_heads\n",
        "        self.attention = layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=embed_dim,dropout=0.1\n",
        "        )\n",
        "        self.dense_proj = keras.Sequential(\n",
        "            [layers.Dense(dense_dim, activation=\"relu\"), layers.Dense(embed_dim),]\n",
        "        )\n",
        "        self.layernorm_1 = layers.LayerNormalization()\n",
        "        self.layernorm_2 = layers.LayerNormalization()\n",
        "        self.supports_masking = True\n",
        "        self.dropout1 = layers.Dropout(0.05)\n",
        "        self.dropout2 = layers.Dropout(0.05)\n",
        "\n",
        "    def call(self, inputs, training, mask=None):\n",
        "        if mask is not None:\n",
        "            padding_mask = tf.cast(mask[:, tf.newaxis, tf.newaxis, :], dtype=\"int32\")\n",
        "        attention_output = self.attention(\n",
        "            query=inputs, value=inputs, key=inputs, attention_mask=padding_mask\n",
        "        )\n",
        "        attention_output = self.dropout1(attention_output, training=training)\n",
        "        proj_input = self.layernorm_1(inputs + attention_output)\n",
        "        proj_output = self.dense_proj(proj_input)\n",
        "        proj_output = self.dropout2(proj_output, training=training)\n",
        "        return self.layernorm_2(proj_input + proj_output)\n",
        "\n",
        "\n",
        "class PositionalEmbedding(layers.Layer):\n",
        "    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):\n",
        "        super(PositionalEmbedding, self).__init__(**kwargs)\n",
        "        self.token_embeddings = layers.Embedding(\n",
        "            input_dim=vocab_size, output_dim=embed_dim\n",
        "        )\n",
        "        self.position_embeddings = layers.Embedding(\n",
        "            input_dim=sequence_length, output_dim=embed_dim\n",
        "        )\n",
        "        self.sequence_length = sequence_length\n",
        "        self.vocab_size = vocab_size\n",
        "        self.embed_dim = embed_dim\n",
        "\n",
        "    def call(self, inputs):\n",
        "        length = tf.shape(inputs)[-1]\n",
        "        positions = tf.range(start=0, limit=length, delta=1)\n",
        "        embedded_tokens = self.token_embeddings(inputs)\n",
        "        embedded_positions = self.position_embeddings(positions)\n",
        "        return embedded_tokens + embedded_positions\n",
        "\n",
        "    def compute_mask(self, inputs, mask=None):\n",
        "        return tf.math.not_equal(inputs, 0)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "reKk7is8YJzk"
      },
      "source": [
        "# IMPORTANT: CHANGE TRAIN TEST SPLIT RATIO HERE"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ratio = 0.6 # change from 0.1 0.2 0.4 0.6 0.8 its not always 0.8 performas better than 0.1"
      ],
      "metadata": {
        "id": "u6AsfBzEPscK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xhHKoqiAUcCo",
        "outputId": "6f0bc6d0-9b08-4cf2-9d00-50908741e423"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4973636 2984182\n",
            "train size proportional to total logs - 0.60000008042406\n"
          ]
        }
      ],
      "source": [
        "train_size = round(df_size * ratio)\n",
        "print(df_size, train_size)\n",
        "print(f\"train size proportional to total logs - {train_size/df_size}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j2t94hDsUcCo",
        "outputId": "79b26374-388f-4b00-f311-5e1a8f28005f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "train ground labels 0 1 counter Counter({0.0: 2686665, 1.0: 400000})\n",
            "test ground labels 0 1 counter Counter({0.0: 1865764, 1.0: 123690})\n",
            "(1989454,) (3086665,) (3086665,) (1989454,)\n"
          ]
        }
      ],
      "source": [
        "from collections import Counter\n",
        "# perform train test split\n",
        "\n",
        "# only take normal logs in target within ratio range \n",
        "data_token_indexed_train = np.append(data_token_indexed[:train_size][true_labels[:train_size]==0], data_token_indexed[df_size:],axis=0)\n",
        "\n",
        "# put remaining to test\n",
        "data_token_indexed_test = data_token_indexed[train_size:df_size]\n",
        "\n",
        "\n",
        "# labels using same split\n",
        "train_ground_labels = np.append(true_labels[:train_size][true_labels[:train_size]==0].flatten(), true_labels[df_size:].flatten(),axis=0)\n",
        "print(f\"train ground labels 0 1 counter {Counter(train_ground_labels)}\")\n",
        "\n",
        "test_ground_labels = true_labels[train_size:df_size] \n",
        "print(f\"test ground labels 0 1 counter {Counter(test_ground_labels)}\")\n",
        "\n",
        "print(data_token_indexed_test.shape, data_token_indexed_train.shape, train_ground_labels.shape, test_ground_labels.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yzjwoT7VUcCp"
      },
      "outputs": [],
      "source": [
        "# verify unique log counts consistent with paper count\n",
        "# Commented because this line takes a long time to run\n",
        "# unique_counter = np.intersect1d(np.unique(data_token_indexed_test), np.unique(data_token_indexed_train), assume_unique=True)\n",
        "# print(f\"Unique log message count - {np.unique(data_token_indexed_test).shape[0] - unique_counter.shape[0]}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_5sxljzWUcCp",
        "outputId": "dbbbb811-ff2d-4ab9-dc80-022d3dc2e2de"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(3086665, 50) (1989454, 50)\n",
            "[1 3 4 5 6 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
            " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
            "(3086665, 50) (3086665,) (1989454, 50) (1989454,)\n",
            "tf.Tensor(\n",
            "[[False False False ...  True  True  True]\n",
            " [False False False ...  True  True  True]\n",
            " [False False False ...  True  True  True]\n",
            " ...\n",
            " [False False False ...  True  True  True]\n",
            " [False False False ...  True  True  True]\n",
            " [False False False ...  True  True  True]], shape=(3086665, 50), dtype=bool) tf.Tensor(\n",
            "[[False False False ...  True  True  True]\n",
            " [False False False ...  True  True  True]\n",
            " [False False False ...  True  True  True]\n",
            " ...\n",
            " [False False  True ...  True  True  True]\n",
            " [False False False ...  True  True  True]\n",
            " [False False False ...  True  True  True]], shape=(1989454, 50), dtype=bool)\n"
          ]
        }
      ],
      "source": [
        "# assign set names\n",
        "x_train = data_token_indexed_train\n",
        "y_train =  train_ground_labels\n",
        "x_test = data_token_indexed_test\n",
        "y_test = test_ground_labels\n",
        "\n",
        "# Add padding to sequences to maxlen 50 according to paper\n",
        "from keras_preprocessing.sequence import pad_sequences\n",
        "\n",
        "x_train = pad_sequences(x_train, maxlen=50, truncating=\"post\", padding=\"post\") \n",
        "x_test = pad_sequences(x_test, maxlen=50, truncating=\"post\", padding=\"post\") \n",
        "print(x_train.shape, x_test.shape)\n",
        "print(x_train[0])\n",
        "print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)\n",
        "\n",
        "## padding masks\n",
        "x_train_masks = tf.equal(x_train, 0)\n",
        "x_test_masks = tf.equal(x_test, 0)\n",
        "print(x_train_masks,x_test_masks)"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ohL4jeccQrWo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F8JtlyvubIxs"
      },
      "source": [
        "# Metrics and model building"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sy2gprEc9yeo",
        "outputId": "c995ca1c-8d9b-4aaa-ddc7-fc98d9553859"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"model_2\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " input_3 (InputLayer)        [(None, 50)]              0         \n",
            "                                                                 \n",
            " positional_embedding_2 (Pos  (None, 50, 16)           38272     \n",
            " itionalEmbedding)                                               \n",
            "                                                                 \n",
            " transformer_encoder_4 (Tran  (None, 50, 16)           2768      \n",
            " sformerEncoder)                                                 \n",
            "                                                                 \n",
            " transformer_encoder_5 (Tran  (None, 50, 16)           2768      \n",
            " sformerEncoder)                                                 \n",
            "                                                                 \n",
            " tf.compat.v1.gather_2 (TFOp  (None, 16)               0         \n",
            " Lambda)                                                         \n",
            "                                                                 \n",
            " lambda_2 (Lambda)           (None,)                   0         \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 43,808\n",
            "Trainable params: 43,808\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n",
            "None\n"
          ]
        }
      ],
      "source": [
        "a = tf.metrics.BinaryAccuracy()\n",
        "def recall_m(y_true, y_pred):\n",
        "    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))\n",
        "    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))\n",
        "    return true_positives / (possible_positives + K.epsilon())\n",
        "\n",
        "def precision_m(y_true, y_pred):\n",
        "    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))\n",
        "    predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))\n",
        "    return true_positives / (predicted_positives + K.epsilon())\n",
        "\n",
        "def f1_m(y_true, y_pred):\n",
        "    precision = precision_m(y_true, y_pred)\n",
        "    recall = recall_m(y_true, y_pred)\n",
        "    return 2*((precision*recall)/(precision+recall+K.epsilon()))\n",
        "\n",
        "def accuracy_m(y_true, y_pred):\n",
        "    a.update_state(y_true,y_pred)\n",
        "    return a.result()\n",
        "\n",
        "def create_model():\n",
        "  embed_dim = 16  # Embedding size for each token\n",
        "  num_heads = 2  # Number of attention heads\n",
        "  ff_dim = 16  # Hidden layer size in feed forward network inside transformer\n",
        "\n",
        "  vocab_size = tokenizer.n_words  # Only consider the top 20k words\n",
        "  maxlen = 50\n",
        "\n",
        "  inputs = layers.Input(shape=(maxlen,))\n",
        "  embedding_layer = PositionalEmbedding(maxlen, vocab_size, embed_dim)\n",
        "  x = embedding_layer(inputs)\n",
        "  transformer_block = TransformerEncoder(embed_dim, num_heads, ff_dim)\n",
        "  transformer_block2 = TransformerEncoder(embed_dim, num_heads, ff_dim)\n",
        "\n",
        "  x = transformer_block(x)\n",
        "  x = transformer_block2(x)\n",
        "  \n",
        "  outputs = tf.gather(x, 0, axis=1) ## take the CLS token as embedding\n",
        "\n",
        "  function = K.sum(K.square(outputs),axis=1)\n",
        "  outputs = tf.keras.layers.Lambda(lambda x: K.sum(K.square(x),axis=1))(outputs)\n",
        "\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
        "  \n",
        "  return model\n",
        "\n",
        "# learning rate decay for optmimizer\n",
        "lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(\n",
        "    initial_learning_rate=0.0001, # 0.0001,\n",
        "    decay_steps=10000,\n",
        "    decay_rate=1-0.001)\n",
        "\n",
        "# optimizer\n",
        "model_opt = tf.keras.optimizers.Adam(learning_rate=lr_schedule,beta_1=0.9, beta_2=0.999)\n",
        "\n",
        "def custom_loss_function(y_true, y_pred):\n",
        "   loss = K.mean((1-y_true)*K.sqrt(y_pred) - (y_true)*K.log(1-K.exp(-K.sqrt(y_pred))))\n",
        "   #loss = K.mean((1-y_true)*K.square(dist) - (y_true)*K.log(1-K.exp(-K.square(dist))))\n",
        "\n",
        "   return loss\n",
        "\n",
        "# Create model with parallel execution - IMPORTANT!!\n",
        "with mirrored_strategy.scope():\n",
        "  model = create_model()\n",
        "\n",
        "# NOTE! the metrics here don't make sense, it requires threshold derivation\n",
        "model.compile(optimizer=model_opt, loss=custom_loss_function, metrics=['binary_accuracy',recall_m,precision_m],loss_weights = [0.2, 1.0])  # this weight is [0.5, 1.0] in paper\n",
        "\n",
        "print(model.summary())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MQmLA2orbjYq"
      },
      "source": [
        "# Evaluation steps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DHXNHOPhOhN6"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import f1_score\n",
        "from sklearn.metrics import confusion_matrix\n",
        "from sklearn.metrics import recall_score\n",
        "from sklearn.metrics import precision_score\n",
        "from sklearn.metrics import accuracy_score\n",
        "from sklearn.metrics import roc_curve\n",
        "from sklearn.metrics import roc_auc_score\n",
        "\n",
        "# We have no idea how original author got the thresholds, we use average of auc score\n",
        "def evaluation_threshold(threshold, y_test, distances):\n",
        "\n",
        "  tr = [threshold]\n",
        "\n",
        "  for t in tr:\n",
        "    res=[]\n",
        "    for d in distances:\n",
        "      if d > t:\n",
        "        res.append(1)\n",
        "      else:\n",
        "        res.append(0)\n",
        "\n",
        "    f1score = f1_score(y_test.astype(np.int32), res)\n",
        "    recall = recall_score(y_test.astype(np.int32), res)\n",
        "    precision = precision_score(y_test.astype(np.int32), res)\n",
        "    accuracy = accuracy_score(y_test.astype(np.int32), res)\n",
        "    print('threshold value: ', t)\n",
        "    print('f1: ', f1score)\n",
        "    print('recall: ',  recall)\n",
        "    print('precision: ', precision)\n",
        "    print('Accuracy: ', accuracy)\n",
        "    print('##################################')\n",
        "    return t, f1score, recall, precision, accuracy\n",
        "\n",
        "max_accuracy = 0\n",
        "max_recall = 0\n",
        "max_precision = 0\n",
        "max_f1score = 0\n",
        "\n",
        "max_accuracy_test = 0\n",
        "max_recall_test = 0\n",
        "max_precision_test = 0\n",
        "max_f1score_test = 0\n",
        "\n",
        "# sometimes it can be a bad run, just rerun to see if results improve\n",
        "for i in range(20):\n",
        "  print(f\"Currently Running Epoch - {i+1}\")\n",
        "  model.fit(x_train, y_train, batch_size=2048, epochs=1)\n",
        "\n",
        "  test_distances = model.predict(x_test,batch_size=2048)\n",
        "  test_auc = roc_auc_score(y_test.astype(np.int32), test_distances)\n",
        "\n",
        "  train_distances = model.predict(x_train,batch_size=2048)\n",
        "  train_auc = roc_auc_score(y_train.astype(np.int32), train_distances)\n",
        "\n",
        "  fpr_test, tpr_test, thresholds_test = roc_curve(y_test.astype(np.int32), test_distances, pos_label=1)\n",
        "  fpr_train, tpr_train, thresholds_train = roc_curve(y_train.astype(np.int32), train_distances, pos_label=1)\n",
        "\n",
        "  print(f'Train AUC: {train_auc}')\n",
        "  print(f'Test AUC: {test_auc}')\n",
        "\n",
        "  if train_auc == 1: # prevent calculation go wrong, not impacting results\n",
        "    train_auc = 0.99\n",
        "  if test_auc == 1:\n",
        "    test_auc = 0.99\n",
        "\n",
        "  print('Using trainset to find threshold:') # this should be legit\n",
        "  t, f1score, recall, precision, accuracy = evaluation_threshold(np.average(thresholds_train[tpr_train>train_auc]), y_test, test_distances)\n",
        "  print('Using testset to find threshold:') # this is what the authors used, which leaks info\n",
        "  t_test, f1score_test, recall_test, precision_test, accuracy_test = evaluation_threshold(np.average(thresholds_test[tpr_test>test_auc]), y_test, test_distances)\n",
        "\n",
        "  if max_accuracy <= accuracy:\n",
        "    max_accuracy = accuracy\n",
        "    max_recall = recall\n",
        "    max_precision = precision\n",
        "    max_f1score = f1score\n",
        "  \n",
        "  if max_accuracy_test <= accuracy_test:\n",
        "    max_accuracy_test = accuracy_test\n",
        "    max_recall_test = recall\n",
        "    max_precision_test = precision_test\n",
        "    max_f1score_test = f1score_test\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"*******\\nauthor's best:\")\n",
        "print(max_accuracy_test, max_precision_test, max_recall_test, max_f1score_test)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y0sJzUS9oLQe",
        "outputId": "27b8c564-a78d-4764-ee36-2558de1e2965"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "*******\n",
            "author's best:\n",
            "0.15477116403762495 0.1388964039888314 1.0 0.24391402677603588\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"*******\\ntrain's best:\")\n",
        "print(max_accuracy, max_precision, max_recall, max_f1score)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6Fmc_HYVoiSm",
        "outputId": "14dd9af3-687d-4678-bda6-e68afe758608"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "*******\n",
            "train's best:\n",
            "0.9210529707608369 0.6332869171352157 1.0 0.7754754054431545\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "TpcmiFCpS6xQ"
      ],
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "gpuClass": "premium",
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}